import sqlparse
from sqlparse.sql import IdentifierList, Identifier
from sqlparse.tokens import Keyword, Whitespace, Newline


def is_valid_sql(sql: str) -> tuple[bool, str]:
    """
    校验SQL是否符合要求
    """
    sql = sql.strip()
    if not sql:
        return False, "SQL为空"

    statements = sqlparse.parse(sql)
    if len(statements) > 1:
        return False, "SQL包含多个语句，只允许单条查询语句"

    parsed = statements[0]
    if parsed.get_type() != 'SELECT':
        return False, "只支持SELECT查询"

    from_seen = False
    table_checked = False

    for token in parsed.tokens:
        if token.ttype in (Whitespace, Newline):
            continue

        if from_seen and not table_checked:
            # 处理标识符列表（多个表名的情况）
            if isinstance(token, IdentifierList):
                for item in token.get_identifiers():
                    table_name = item.get_real_name() or str(item).split()[0]
                    if table_name.lower() != "sales":
                        return False, f"只能查询sales表，不能查询{table_name}"
                table_checked = True
            # 处理单个标识符（单个表名的情况）
            elif isinstance(token, Identifier):
                table_name = token.get_real_name() or str(token).split()[0]
                if table_name.lower() != "sales":
                    return False, f"只能查询sales表，不能查询{table_name}"
                table_checked = True
            # 处理其他情况
            else:
                table_name = str(token).split()[0].strip()
                if table_name.lower() != "sales":
                    return False, f"只能查询sales表，不能查询{table_name}"
                table_checked = True

        if token.ttype is Keyword and token.value.upper() == "FROM":
            from_seen = True

    if not from_seen:
        return False, "SQL缺少FROM子句"

    if not table_checked:
        return False, "未找到有效的表名"

    return True, "SQL格式有效"


if __name__ == "__main__":
    test_cases = [
        # 合法场景
        ("SELECT * FROM sales", (True, "SQL格式有效")),
        ("select name from sales as t", (True, "SQL格式有效")),  # 小写+别名

        # 非法场景：表名错误
        ("SELECT * FROM other_table", (False, "只能查询sales表，不能查询other_table")),

        # 非法场景：非SELECT语句
        ("UPDATE data_sales SET x=1", (False, "只支持SELECT查询")),

        # 非法场景：多条语句
        ("SELECT * FROM data_sales; SELECT 1", (False, "SQL包含多个语句，只允许单条查询语句")),

        # 非法场景：缺少FROM
        ("SELECT 123", (False, "SQL缺少FROM子句")),

        # 边界场景：空SQL
        ("   ", (False, "SQL为空")),
    ]

    # 执行测试
    for i, (sql, expected) in enumerate(test_cases, 1):
        result = is_valid_sql(sql)
        if result == expected:
            print(f"用例{i}：通过 ✅")
        else:
            print(f"用例{i}：失败 ❌\nSQL: {sql}\n实际: {result}\n预期: {expected}\n---")